# !/usr/bin/env python
# encoding: utf-8


"""
  @author: gaogao
  @file: align_test.py
  @time: 2022/7/27 14:45
  @desc:
"""
from cProfile import label
import time
import numpy as np
import pyvisa as visa
from configparser import ConfigParser
from matplotlib import pyplot as plt
import socket

inst_light = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
inst_light.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)  # 设置缓冲区
inst_light.connect(("192.168.1.111", 5000))
cmd = f"<OSW01_OUT_01_05>"
inst_light.send(cmd.encode())
time.sleep(1)

n7744_address = "TCPIP0::192.168.1.11::5025::SOCKET"
n7786_address = "TCPIP0::192.168.1.21::5025::SOCKET"
n8164_address = "TCPIP0::192.168.1.20::5025::SOCKET"
n7744_inst = visa.ResourceManager().open_resource(n7744_address)
n7786_inst = visa.ResourceManager().open_resource(n7786_address)
n8164_inst = visa.ResourceManager().open_resource(n8164_address)
n8164_inst.timeout = 6000
n8164_inst.read_termination = '\n'
n8164_inst.write_termination = '\n'

n7744_inst.timeout = 6000
n7744_inst.read_termination = '\n'
n7744_inst.write_termination = '\n'
n7786_inst.timeout = 6000
n7786_inst.read_termination = '\n'
n7786_inst.write_termination = '\n'


def align():
    n8164_inst.write("sour0:wav 1330NM")
    n8164_inst.write("sour0:pow:unit 0")
    n8164_inst.write("sour0:pow 10 dbm")
    time.sleep(4)
    n8164_inst.write(":OUTPUT:CHANNEL:STATE 1")

    n7744_inst.write("SYST:PRES;*OPC?")
    n7744_inst.write("SENS1:FUNC:STAT LOGG,STOP")
    n7744_inst.write("SENS1:POW:RANG 10DBM")
    n7744_inst.write("SENS1:POW:GAIN:AUTO 0")
    n7744_inst.write("TRIG1:INP SME")
    n7744_inst.write("SENS1:POW:WAV 1330nm")
    n7744_inst.write("SENS1:FUNC:PAR:LOGG 20000,100US")

    n7786_inst.write("SYST:PRES;*OPC?")
    n7786_inst.write("STAB:STAB 0")

    n7786_inst.write("PCON:SEQ:DCOM 1")
    n7786_inst.write("PCON:SEQ:RRAT 8")

    # n7786_inst.query(f"PCON:GEN:RAND? 20000,1000")
    n7786_inst.query(f"PCON:LOAD:ARB?")
    # n7786_inst.query(f"PCON:SAVE:ARB?")
    n7786_inst.write("PCON:SEQ:HOLD 640")  # Define an external trigger delay after changing the SOP value
    n7786_inst.write("PCON:SEQ:SMOD 3")
    n7786_inst.write("PCON:REP 0")
    n7786_inst.write(":TRIG:CONF 4")
    n7786_inst.write("POL:TRIG:INP SME")
    n7786_inst.write(f"POL:SWE:SAMP 20000")
    n7786_inst.write(f"POL:SWE:SRAT 8000HZ,100US")
    n7786_inst.write("POL:SWE:LOOP 0")
    n7786_inst.write("POL:GAIN 4")

    n7786_inst.write(f"POL:WAV 1330NM")
    n7786_inst.write(f"POL:SWE:STEP 0")
    # n7786_inst.write(":POL:SWE:START SOP")
    n7786_inst.write(f"POL:SWE:START")
    n7744_inst.write("SENS1:FUNC:STAT LOGG,STAR")

    n7786_inst.write(f"PCON:STAR")
    time.sleep(3)
    # n7786_inst.write(f"SENS1:FUNC:STAT LOGG,START")
    n7786_inst.write(f"TRIG 1")
    n7744_inst.write(":TRIG 2")

    while True:
        state = n7744_inst.query("SENS1:FUNC:STAT?")
        if state == "LOGGING_STABILITY,COMPLETE":
            result_list = n7744_inst.query_binary_values("SENS1:FUNC:RES?")
            break

    max_index = result_list.index(max(result_list))
    min_index = result_list.index(min(result_list))

    while True:
        log_value = n7786_inst.query("POL:SWE:SAMP:CURR?")
        print(log_value)
        if log_value == "+0":
            polarization_data = n7786_inst.query_binary_values("POL:SWE:GET? NORM")
            break

    n7786_inst.write("POL:STOP")  # Stops logging
    n7786_inst.write("PCON:STOP")
    state = n7786_inst.query("*OPC?")

    if state == "1":
        n7786_inst.write("PCON:SEQ:HOLD 1")
        n7786_inst.write(":TRIG:CONF 2")
        n7786_inst.write("STAB:STAB 1")

    max_polar = [polarization_data[max_index * 3], polarization_data[max_index * 3 + 1],
                 polarization_data[max_index * 3 + 2]]
    min_polar = [polarization_data[min_index * 3], polarization_data[min_index * 3 + 1],
                 polarization_data[min_index * 3 + 2]]
    n7786_inst.write(f"STAB:SOP {max_polar[0]},{max_polar[1]},{max_polar[2]}")


def get_refrence_data(ini_path=r"C:\Users\Jingw\PycharmProjects\HS_ATS_VS\static\refrence_data.ini",
                      section="loss_sys"):
    config = ConfigParser()
    config.read(ini_path, encoding='utf-8-sig')
    values = config.items(section)
    xy_data_tuple = (eval(values[0][1]), eval(values[1][1]))
    return xy_data_tuple


def get_wave_power_value(wave, refrence_data):
    wave = float(wave)
    if wave in refrence_data[0]:
        wave_index = refrence_data[0].index(wave)
    else:
        wave = min(refrence_data[0], key=lambda x: abs(x - wave))
        wave_index = refrence_data[0].index(wave)
    return refrence_data[1][wave_index]


def get_step_list(start_value, stop_value, step_value):
    if float(stop_value) > float(start_value):
        points_num = abs((float(stop_value) - float(start_value)) / (
            float(step_value))) + 1
        data_list = [round(float(start_value) + i * float(step_value), 3) for i in range(int(points_num))]
        if data_list:
            return data_list
    elif float(stop_value) < float(start_value):
        points_num = abs((float(stop_value) - float(start_value)) / (
            float(step_value))) + 1
        data_list = [round(float(start_value) - i * float(step_value), 3) for i in range(int(points_num))]
        if data_list:
            return data_list
    else:
        data_list = [float(start_value)]
        return data_list


def n7744_auto_range(meas_data, current_range):
    range_list = ["10", "0", "-10", "-20", "-30"]
    current_index = range_list.index(str(int(float(current_range))))
    max_value = max(meas_data)
    print("current_index", current_index)
    print("max_value", max_value)
    is_variety = False
    if float(max_value) >= float(current_range) * 0.9:
        if current_index > 0:
            current_range = range_list[current_index - 1]
            is_variety = True
    elif float(max_value) < float(range_list[current_index + 1]) * 0.9:
        if current_index < len(range_list)-1:
            current_range = range_list[current_index + 1]
            is_variety = True
    return is_variety, current_range


def p_λ():
    n7744_inst.write("SENS1:FUNC:STAT LOGG,STOP")
    x_data = get_step_list(1260, 1360, 1)
    points = len(x_data)
    print("points", points)
    n8164_inst.write("sour0:pow:unit 0")
    n8164_inst.write("sour0:pow 10 dbm")
    time.sleep(2)
    p_in = n8164_inst.query("sour0:pow?")
    n8164_inst.write("WAVelength:SWEep:MODE CONT")
    n8164_inst.write("wav:swe:rep ONEW")
    n8164_inst.write("wav:swe:spe 100nm/s")
    n8164_inst.write(f"WAVelength:SWEep:STARt 1260NM")
    n8164_inst.write(f"WAVelength:SWEep:STOP 1360NM")
    n8164_inst.write("WAVelength:SWEep:STEP 1NM")

    n8164_inst.write("trig:conf 1")
    n8164_inst.write(f":TRIGger:OUTPut STF")
    n8164_inst.write(f":TRIGger:INPut SWS")
    n8164_inst.write(":OUTPUT:CHANNEL:STATE 1")
    n8164_inst.write(f"wav:swe 1")
    time.sleep(2)
    n7744_inst.query("SYST:PRES;*OPC?")
    n7744_inst.write(f":sens1:pow:unit 0")
    n7744_inst.write(f"SENS1:POW:RANG 10DBM")
    n7744_inst.write(f"SENS1:POW:GAIN:AUTO 0")
    n7744_inst.write(f":sens1:pow:atim 100us")
    n7744_inst.write(f"TRIG1:INP SME")
    n7744_inst.write(f"SENS1:POW:WAV 1310nm")
    n7744_inst.write(f"SENS1:FUNC:PAR:LOGG {points},100US")
    n7744_inst.write(f"SENS1:FUNC:STAT LOGG,STAR")
    time.sleep(3)
    n7744_inst.write(f":TRIG 2")
    time.sleep(1)
    while True:
        state = n7744_inst.query("SENS1:FUNC:STAT?").strip()
        if state == "LOGGING_STABILITY,COMPLETE":
            y_data = n7744_inst.query_binary_values("SENS1:FUNC:RES?")
            break

    n8164_inst.write("wav:swe STOP")
    y_data = [float(i) for i in y_data]
    wave_power_tuple = get_refrence_data()
    for wave_index in range(len(x_data)):
        loss_sys = get_wave_power_value(x_data[wave_index],
                                        wave_power_tuple)
        y_data[wave_index] = float(p_in) - y_data[wave_index] - loss_sys
    min_power = min(y_data)
    min_index = y_data.index(min_power)
    wave = x_data[min_index]
    loss_gc_list = [float(i) / 2 for i in y_data]
    return x_data, loss_gc_list


def p_λ1(start_wave, stop_wave):
    current_range = -10
    n7744_inst.write("SENS1:FUNC:STAT LOGG,STOP")
    n7744_inst.query("SYST:PRES;*OPC?")
    while True:
        
        x_data = get_step_list(int(start_wave), int(stop_wave), 1)
        points = len(x_data)
        print("points", points)
        n8164_inst.write("sour0:pow:unit 0")
        n8164_inst.write("sour0:pow 10 dbm")
        time.sleep(1)
        p_in = n8164_inst.query("sour0:pow?")
        n8164_inst.write("WAVelength:SWEep:MODE CONT")
        n8164_inst.write("wav:swe:rep ONEW")
        n8164_inst.write("wav:swe:spe 100nm/s")
        n8164_inst.write(f"WAVelength:SWEep:STARt {start_wave}NM")
        n8164_inst.write(f"WAVelength:SWEep:STOP {stop_wave}NM")
        n8164_inst.write("WAVelength:SWEep:STEP 1NM")

        n8164_inst.write("trig:conf 1")
        n8164_inst.write(f":TRIGger:OUTPut STF")
        n8164_inst.write(f":TRIGger:INPut SWS")
        n8164_inst.write(":OUTPUT:CHANNEL:STATE 1")
        n8164_inst.write(f"wav:swe 1")
        time.sleep(1)
        
        n7744_inst.write(f":sens1:pow:unit 0")
        n7744_inst.write(f"SENS1:POW:RANG {current_range}DBM")
        n7744_inst.write(f"SENS1:POW:GAIN:AUTO 0")
        n7744_inst.write(f":sens1:pow:atim 100us")
        n7744_inst.write(f"TRIG1:INP SME")
        n7744_inst.write(f"SENS1:POW:WAV 1310nm")
        n7744_inst.write(f"SENS1:FUNC:PAR:LOGG {points},100US")
        n7744_inst.write(f"SENS1:FUNC:STAT LOGG,STAR")
        time.sleep(2)
        n7744_inst.write(f":TRIG 2")
        time.sleep(1)
        while True:
            state = n7744_inst.query("SENS1:FUNC:STAT?").strip()
            if state == "LOGGING_STABILITY,COMPLETE":
                y_data = n7744_inst.query_binary_values("SENS1:FUNC:RES?")
                break

        n8164_inst.write("wav:swe STOP")
        y_data = [float(i) for i in y_data]
        current_range = n7744_inst.query("SENS1:POW:RANG?")
        is_verify, handle_range = n7744_auto_range(y_data, current_range)
        print(is_verify)
        print(handle_range)
        if is_verify:
            current_range = handle_range
            time.sleep(15)
            continue
        else:
            break
    
    wave_power_tuple = get_refrence_data()
    for wave_index in range(len(x_data)):
        loss_sys = get_wave_power_value(x_data[wave_index],
                                        wave_power_tuple)
        y_data[wave_index] = float(p_in) - y_data[wave_index] - loss_sys
    min_power = min(y_data)
    min_index = y_data.index(min_power)
    wave = x_data[min_index]
    loss_gc_list = [float(i) / 2 for i in y_data]
    return x_data, loss_gc_list


def p_λ2():
    # P-拉姆达
    n7744_inst.write("SENS1:FUNC:STAT LOGG,STOP")
    n8164_inst.write("sour0:pow:unit 0")
    n8164_inst.write("sour0:pow 10 dbm")

    time.sleep(3)
    n8164_inst.write(":OUTPUT:CHANNEL:STATE 1")
    p_in = n8164_inst.query("sour0:pow?")
    wave_power_tuple = get_refrence_data()
    power_list = []
    wave_list = get_step_list(1260, 1360, 1)
    wave_list.reverse()
    n7744_inst.query("SYST:PRES;*OPC?")
    n7744_inst.write(":sens1:pow:unit 0")
    n7744_inst.write(":SENSe1:POWer:RANGe:AUTO 1")
    # n7744_inst.write(f"SENS1:POW:WAV 1310nm")
    n7744_inst.write(":sens1:pow:atim 100us")
    for wave in wave_list:
        print("wave", wave)
        n8164_inst.write(f"sour0:wav {wave}NM")
        loss_sys = get_wave_power_value(wave, wave_power_tuple)
        if wave == "1260":
            time.sleep(2)
        else:
            time.sleep(0.5)
        power = float(n7744_inst.query(":FETCH1:POW?"))
        power_value = float(p_in) \
                      - power \
                      - loss_sys
        power_list.append(power_value)
    print("power_list", power_list)
    loss_gc_list = [float(i) / 2 for i in power_list]

    return wave_list, loss_gc_list


def plot_curve(curve_list):
    plt.figure(figsize=(20, 20), dpi=100)
    fig, ax = plt.subplots()
    ax.set_title(f"Loss GC curve")
    for i in range(len(curve_list)):
        ax.plot(curve_list[i][0], curve_list[i][1], label=str(i))
        ax.set_xlabel("WaveLength")
        ax.set_ylabel("Power")
        ax.grid(True)
        ax.legend()
    ax.invert_yaxis()
    plt.show()


if __name__ == "__main__":


    align()
    curve_list = []
    x2_list,y2_list = p_λ()
    curve_list.append([x2_list, y2_list])
    x_list, y_list = p_λ1(1260,1310)
    x1_list, y1_list = p_λ1(1311,1360)
    x_list.extend(x1_list)
    y_list.extend(y1_list)
    curve_list.append([x_list, y_list])
    time.sleep(10)
    plot_curve(curve_list)
